from flask import Flask
from flask import render_template
from flask import request
from flask import redirect
import torch
import cv2
import numpy as np
from NumberNet import NumberNet
from flask_cors import CORS

model = torch.load("./static/models/numbers_model1.pth", weights_only=False)

app = Flask(__name__)
CORS(app)


# ---- 渲染整个识别的网页-----
@app.route("/")
def detect_number():
    return render_template("index.html")


# ---- 上传文件，并识别该内容 ----
@app.route('/upload', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        f = request.files['the_file']  # 获取文件对象
        f.save('./static/upload/uploaded_file.png')  # 将文件对象保存到服务器
        # ---- 识别本地文件--------
        img = cv2.imread("./static/upload/uploaded_file.png")
        print(img.shape)
        # img = cv2.resize(img,(28,28))
        img = np.expand_dims(img, 0)
        img = torch.from_numpy(img)
        img = torch.permute(img, [0, 3, 1, 2])
        print(img.shape)
        model_cpu = model.to(torch.device("cpu"))
        predict = model_cpu(img.float())
        result = torch.argmax(predict, dim=-1)

    return {"code": 202, "result": result[0].item()}


if __name__ == "__main__":
    app.run("localhost", 9000, debug=True)
